from typing import Any
from transformers import BertTokenizer, BertModel


from qdrant_client import QdrantClient

tokenizer = BertTokenizer.from_pretrained("D:/bert-base-chinese")
model = BertModel.from_pretrained("D:/bert-base-chinese")

collection_name = "webui"


def EmbeddingText(text: str) -> Any:
    inputs = tokenizer(text, return_tensors="pt")
    outputs = model(**inputs)
    return outputs.last_hidden_state[:, 0, :].tolist()[0]


# 连接数据库
client = QdrantClient(host="10.254.32.20", port=6333)
collection_name = "webui"

vector = EmbeddingText("饼图")
results = client.query_points(
    collection_name=collection_name, query=vector, score_threshold=0.6, limit=30
)

print(type(results))
print(results)


for hit in results.points:
    print(hit.payload)  # 例如提取 'desc' 字段
